import os
import numpy as np
import matplotlib.pyplot as plt
from scipy.io import loadmat
from typing import Dict, Any, Tuple

def save_raw_segment(args):
    """保存原始信号片段，与预处理脚本保持一致的保存格式"""
    start, window_length, raw_signal, save_path, file_id, segment_idx = args
    segment = raw_signal[start : start + window_length]
    if len(segment) == window_length:
        sample = {
            'ecg': segment,
            'file_id': file_id,
            'segment_idx': segment_idx
        }
        np.save(save_path, sample)
        return True
    return False

def process_raw_file(file_name: str, data_dir: str, save_dir: str, split_name: str, fs: float = 512.0) -> None:
    """处理原始文件并保存切片"""
    print(f'处理原始文件: {file_name}')
    file_path = os.path.join(data_dir, file_name)
    data = loadmat(file_path)
    raw_signal = data.get('segmentData')
    
    if raw_signal is None:
        print(f'文件 {file_name} 中未找到变量 "segmentData"')
        return

    raw_signal = np.squeeze(raw_signal)
    
    # 获取文件ID
    full_sample_id = file_name.split('.')[0]  # 例如: "001-1"
    
    # 只保存第一个片段用于对比
    window_length = int(fs * 5)  # 5秒窗口
    start = 0
    
    if split_name == 'predict_data':
        save_path = os.path.join(save_dir, f'{full_sample_id}_segment_0.npy')
    else:
        save_path = os.path.join(save_dir, f'{full_sample_id}_sample_0.npy')
    
    save_raw_segment((start, window_length, raw_signal, save_path, full_sample_id, 0))

def load_signals(raw_dir: str, processed_dir: str, file_name: str, split_name: str) -> Tuple[np.ndarray, np.ndarray]:
    """加载原始信号和处理后的信号"""
    # 获取完整的样本ID
    full_sample_id = file_name.split('.')[0]  # 例如: "001-1"
    
    # 加载原始信号片段
    raw_path = os.path.join(raw_dir, f'{full_sample_id}_raw_sample_0.npy')
    raw_data = np.load(raw_path, allow_pickle=True).item()
    raw_segment = raw_data['ecg']
    
    # 加载处理后的信号片段 - 根据split_name使用不同的命名格式
    if split_name == 'predict_data':
        proc_path = os.path.join(processed_dir, f'{full_sample_id}_segment_0.npy')
    else:
        proc_path = os.path.join(processed_dir, f'{full_sample_id}_sample_0.npy')
    
    proc_data = np.load(proc_path, allow_pickle=True).item()
    proc_segment = proc_data['ecg']
    
    return raw_segment, proc_segment

def plot_signal_comparison(raw_segment: np.ndarray, processed_segment: np.ndarray, 
                         file_id: str, fs: float = 512.0, save_dir: str = 'results/signal_comparison'):
    """绘制对比图"""
    time = np.arange(len(raw_segment)) / fs

    plt.figure(figsize=(15, 8))
    
    # 绘制原始信号
    plt.subplot(211)
    plt.plot(time, raw_segment, 'b-', label='原始信号')
    plt.grid(True)
    plt.legend()
    plt.title(f'样本 {file_id} 的信号对比')
    plt.ylabel('幅值')
    
    # 绘制处理后的信号
    plt.subplot(212)
    plt.plot(time, processed_segment, 'r-', label='滤波后信号')
    plt.grid(True)
    plt.legend()
    plt.xlabel('时间 (秒)')
    plt.ylabel('幅值')
    
    os.makedirs(save_dir, exist_ok=True)
    plt.savefig(os.path.join(save_dir, f'{file_id}_comparison.png'))
    plt.close()

def compare_signals(data_dir: str, processed_dir: str, raw_save_dir: str, split_name: str):
    """比较原始信号和处理后的信号"""
    os.makedirs(raw_save_dir, exist_ok=True)
    
    # 首先处理所有原始文件
    print("处理原始文件...")
    mat_files = [f for f in os.listdir(data_dir) if f.endswith('.mat')]
    for file_name in mat_files:
        process_raw_file(file_name, data_dir, raw_save_dir, split_name)
    
    # 然后进行信号对比
    print("\n生成对比图...")
    for file_name in mat_files:
        try:
            raw_segment, proc_segment = load_signals(raw_save_dir, processed_dir, file_name, split_name)
            # 使用完整的样本ID作为标识
            full_sample_id = file_name.split('.')[0]
            plot_signal_comparison(raw_segment, proc_segment, full_sample_id)
            print(f"完成样本 {full_sample_id} 的信号对比图绘制")
        except FileNotFoundError as fnf_error:
            print(f"处理 {file_name} 时出错: {fnf_error}")
        except Exception as e:
            print(f"处理 {file_name} 时发生其他错误: {e}")

if __name__ == '__main__':
    data_dirs = {
        'train_data': 'data/train_data',
        'test_data': 'data/test_data',
        'predict_data': 'data/predict_data'  # 添加预测数据目录
    }
    processed_dir = 'data/processed'
    raw_save_dir = 'data/raw_segments'
    
    for split_name, data_dir in data_dirs.items():
        print(f"\n处理{split_name}...")
        split_processed_dir = os.path.join(processed_dir, split_name)
        split_raw_dir = os.path.join(raw_save_dir, split_name)
        
        if os.path.exists(split_processed_dir):
            compare_signals(data_dir, split_processed_dir, split_raw_dir, split_name)
        else:
            print(f"目录不存在: {split_processed_dir}")
